package com.zong.ai.agent.chatMemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.util.DefaultInstantiatorStrategy;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * @author: zongzi
 * @description: 基于文件的"会话记忆"，使用Kryo序列化
 * @date: 2025/6/18
 */
public class InFileChatMemory implements ChatMemory {

    private static final String FILE_PREFIX = "you-gi-on-";
    private static final String FILE_SUFFIX = ".kryo";
    private static final String DEFAULT_DIR = "chat-memory";
    
    private final Kryo kryo;
    private final String memoryDir;

    public InFileChatMemory() {
        this(DEFAULT_DIR);
    }

    public InFileChatMemory(String memoryDir) {
        this.memoryDir = memoryDir;
        this.kryo = new Kryo();
        
        // 配置 Kryo 的实例化策略，支持没有无参构造函数的类
        kryo.setInstantiatorStrategy(new DefaultInstantiatorStrategy(new StdInstantiatorStrategy()));
        
        // 注册需要序列化的类
        kryo.setRegistrationRequired(false);
        
        // 确保目录存在
        File dir = new File(memoryDir);
        if (!dir.exists()) {
            dir.mkdirs();
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        if (conversationId == null || messages == null || messages.isEmpty()) {
            return;
        }

        // 从文件中获取现有消息
        List<Message> existingMessages = loadFromFile(conversationId);
        
        // 添加新消息
        existingMessages.addAll(messages);
        
        // 保存到文件
        saveToFile(conversationId, existingMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        if (conversationId == null) {
            return new ArrayList<>();
        }

        List<Message> messages = loadFromFile(conversationId);
        
        if (messages.isEmpty()) {
            return new ArrayList<>();
        }

        // 返回最后N条消息
        if (lastN <= 0) {
            return new ArrayList<>(messages);
        }
        
        int start = Math.max(0, messages.size() - lastN);
        return new ArrayList<>(messages.subList(start, messages.size()));
    }

    @Override
    public void clear(String conversationId) {
        if (conversationId == null) {
            return;
        }

        // 删除文件
        File file = new File(memoryDir, getFileName(conversationId));
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 生成文件名
     */
    private String getFileName(String conversationId) {
        return FILE_PREFIX + conversationId + FILE_SUFFIX;
    }

    /**
     * 从文件加载消息列表
     */
    private List<Message> loadFromFile(String conversationId) {
        File file = new File(memoryDir, getFileName(conversationId));
        
        if (!file.exists()) {
            return new ArrayList<>();
        }

        try (Input input = new Input(new FileInputStream(file))) {
            @SuppressWarnings("unchecked")
            List<Message> messages = kryo.readObject(input, ArrayList.class);
            return messages != null ? messages : new ArrayList<>();
        } catch (FileNotFoundException e) {
            return new ArrayList<>();
        } catch (Exception e) {
            // 如果反序列化失败，返回空列表并删除损坏的文件
            System.err.println("Failed to load chat memory from file: " + file.getAbsolutePath() + ", error: " + e.getMessage());
            file.delete();
            return new ArrayList<>();
        }
    }

    /**
     * 保存消息列表到文件
     */
    private void saveToFile(String conversationId, List<Message> messages) {
        File file = new File(memoryDir, getFileName(conversationId));
        
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messages);
        } catch (Exception e) {
            System.err.println("Failed to save chat memory to file: " + file.getAbsolutePath() + ", error: " + e.getMessage());
        }
    }
}
